package com.wanma.Collection;

import java.util.Deque;
import java.util.HashSet;
import java.util.Iterator;
import java.util.LinkedList;
import java.util.List;
import java.util.Set;
import java.util.function.Function;

public class WeekTree {

	public static <T> Iterator<T> walkTree(T root, Function<T, List<T>> getChildren, Function<Deque<T>, T> pollFunction, Function<Deque<T>, T> peekFunction) {
		   return new Iterator<T>() {
		      private T next = root;
		      private Deque<T> deque = new LinkedList<>();
		      private Set<T> visited = new HashSet<>();

		      public boolean hasNext() {
		         return next != null;
		      }

		      public T next() {
		         deque.push(next);
		         visited.add(next);
		         T current = next;
		         for (next = null; !deque.isEmpty(); pollFunction.apply(deque)) {
		            next = getChildren.apply(peekFunction.apply(deque)).stream().filter(n -> !visited.contains(n)).findFirst().orElse(null);
		            if (next != null) {
		               break;
		            }
		         }
		         return current;
		      }
		   };
		}
}
